#!/usr/bin/env python3
"""
Dolphin配置验证工具
检查模型文件、配置文件等是否正确设置
"""

import os
import sys
from pathlib import Path
from typing import List, Tuple, Dict, Any

try:
    from omegaconf import OmegaConf
    import torch
    from PIL import Image
    import cv2
    import numpy as np
except ImportError as e:
    print(f"❌ 缺少必要的依赖包: {e}")
    print("请运行: pip install -r requirements.txt")
    sys.exit(1)


class ConfigValidator:
    def __init__(self):
        self.issues = []
        self.warnings = []
        
    def log_issue(self, message: str):
        """记录问题"""
        self.issues.append(message)
        print(f"❌ 错误: {message}")
    
    def log_warning(self, message: str):
        """记录警告"""
        self.warnings.append(message)
        print(f"⚠️  警告: {message}")
    
    def log_success(self, message: str):
        """记录成功"""
        print(f"✅ {message}")
    
    def check_python_version(self) -> bool:
        """检查Python版本"""
        print("🐍 检查Python版本...")
        
        version = sys.version_info
        if version.major == 3 and version.minor >= 8:
            self.log_success(f"Python版本: {version.major}.{version.minor}.{version.micro}")
            return True
        else:
            self.log_issue(f"Python版本过低: {version.major}.{version.minor}.{version.micro}，需要Python 3.8+")
            return False
    
    def check_dependencies(self) -> bool:
        """检查依赖包"""
        print("\n📦 检查Python依赖包...")
        
        required_packages = [
            ('torch', '深度学习框架'),
            ('torchvision', '计算机视觉库'),
            ('transformers', 'Transformer模型库'),
            ('PIL', '图像处理库'),
            ('cv2', 'OpenCV计算机视觉库'),
            ('numpy', '科学计算库'),
            ('omegaconf', '配置管理库'),
            ('fastapi', 'Web框架'),
            ('uvicorn', 'ASGI服务器'),
            ('pydantic', '数据验证库'),
            ('requests', 'HTTP客户端库'),
            ('pymupdf', 'PDF处理库')
        ]
        
        missing_packages = []
        
        for package, description in required_packages:
            try:
                __import__(package)
                self.log_success(f"{package} - {description}")
            except ImportError:
                missing_packages.append(package)
                self.log_issue(f"缺少依赖包: {package} - {description}")
        
        if missing_packages:
            print(f"\n安装缺少的依赖包:")
            print(f"pip install {' '.join(missing_packages)}")
            return False
        
        return True
    
    def check_gpu_support(self) -> bool:
        """检查GPU支持"""
        print("\n🎮 检查GPU支持...")
        
        if torch.cuda.is_available():
            gpu_count = torch.cuda.device_count()
            for i in range(gpu_count):
                gpu_name = torch.cuda.get_device_name(i)
                gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
                self.log_success(f"GPU {i}: {gpu_name} ({gpu_memory:.1f}GB)")
            
            # 测试GPU内存分配
            try:
                test_tensor = torch.randn(100, 100).cuda()
                del test_tensor
                torch.cuda.empty_cache()
                self.log_success("GPU内存分配测试通过")
            except Exception as e:
                self.log_warning(f"GPU内存分配测试失败: {e}")
            
            return True
        else:
            self.log_warning("未检测到CUDA支持，将使用CPU模式")
            return False
    
    def check_config_file(self, config_path: str = "config/Dolphin.yaml") -> bool:
        """检查配置文件"""
        print(f"\n⚙️  检查配置文件: {config_path}")
        
        if not os.path.exists(config_path):
            self.log_issue(f"配置文件不存在: {config_path}")
            return False
        
        try:
            config = OmegaConf.load(config_path)
            self.log_success("配置文件加载成功")
            
            # 检查必要的配置项
            required_keys = [
                'model.model_name_or_path',
                'model.tokenizer_path',
                'model.max_length',
                'model.decoder_layer',
                'model.hidden_dimension',
                'model.swin_args'
            ]
            
            for key in required_keys:
                if not OmegaConf.select(config, key):
                    self.log_issue(f"配置文件缺少必要项: {key}")
                    return False
                else:
                    value = OmegaConf.select(config, key)
                    self.log_success(f"{key}: {value}")
            
            return True
            
        except Exception as e:
            self.log_issue(f"配置文件解析失败: {e}")
            return False
    
    def check_model_files(self, config_path: str = "config/Dolphin.yaml") -> bool:
        """检查模型文件"""
        print("\n🤖 检查模型文件...")
        
        if not os.path.exists(config_path):
            self.log_issue(f"配置文件不存在，无法检查模型文件")
            return False
        
        try:
            config = OmegaConf.load(config_path)
            
            # 检查模型权重文件
            model_path = config.model.model_name_or_path
            if not os.path.exists(model_path):
                self.log_issue(f"模型文件不存在: {model_path}")
                return False
            else:
                size_mb = os.path.getsize(model_path) / 1024 / 1024
                self.log_success(f"模型文件: {model_path} ({size_mb:.1f}MB)")
            
            # 检查分词器文件
            tokenizer_path = config.model.tokenizer_path
            if not os.path.exists(tokenizer_path):
                self.log_issue(f"分词器文件不存在: {tokenizer_path}")
                return False
            else:
                size_kb = os.path.getsize(tokenizer_path) / 1024
                self.log_success(f"分词器文件: {tokenizer_path} ({size_kb:.1f}KB)")
            
            return True
            
        except Exception as e:
            self.log_issue(f"检查模型文件时出错: {e}")
            return False
    
    def check_directories(self) -> bool:
        """检查必要的目录结构"""
        print("\n📁 检查目录结构...")
        
        required_dirs = [
            'config',
            'checkpoints',
            'utils',
            'demo'
        ]
        
        for dir_name in required_dirs:
            if os.path.exists(dir_name):
                self.log_success(f"目录存在: {dir_name}")
            else:
                self.log_warning(f"目录不存在: {dir_name}")
        
        # 检查输出目录是否可创建
        test_dirs = ['temp_uploads', 'api_results']
        for dir_name in test_dirs:
            try:
                os.makedirs(dir_name, exist_ok=True)
                self.log_success(f"输出目录可创建: {dir_name}")
            except Exception as e:
                self.log_issue(f"无法创建目录 {dir_name}: {e}")
        
        return True
    
    def test_model_loading(self, config_path: str = "config/Dolphin.yaml") -> bool:
        """测试模型加载"""
        print("\n🔄 测试模型加载...")
        
        try:
            # 尝试导入和初始化模型（简化版本）
            from chat import DOLPHIN
            config = OmegaConf.load(config_path)
            
            # 注意：这里只是测试能否实例化，不进行实际推理
            self.log_success("模型类导入成功")
            
            # 检查模型文件是否可以被torch加载
            model_path = config.model.model_name_or_path
            try:
                checkpoint = torch.load(model_path, map_location='cpu')
                self.log_success(f"模型文件可以被torch加载，包含 {len(checkpoint)} 个键")
                del checkpoint  # 释放内存
            except Exception as e:
                self.log_issue(f"模型文件加载失败: {e}")
                return False
            
            return True
            
        except Exception as e:
            self.log_issue(f"模型加载测试失败: {e}")
            return False
    
    def check_demo_files(self) -> bool:
        """检查示例文件"""
        print("\n🖼️  检查示例文件...")
        
        demo_dirs = ['demo/page_imgs', 'demo/element_imgs']
        found_files = 0
        
        for demo_dir in demo_dirs:
            if os.path.exists(demo_dir):
                files = list(Path(demo_dir).glob('*'))
                image_files = [f for f in files if f.suffix.lower() in ['.jpg', '.jpeg', '.png', '.pdf']]
                
                if image_files:
                    self.log_success(f"{demo_dir}: 找到 {len(image_files)} 个示例文件")
                    found_files += len(image_files)
                else:
                    self.log_warning(f"{demo_dir}: 没有找到示例文件")
            else:
                self.log_warning(f"示例目录不存在: {demo_dir}")
        
        if found_files > 0:
            self.log_success(f"总计找到 {found_files} 个示例文件用于测试")
        else:
            self.log_warning("没有找到示例文件，建议添加一些测试图像")
        
        return True
    
    def run_validation(self, config_path: str = "config/Dolphin.yaml") -> Dict[str, Any]:
        """运行完整的配置验证"""
        print("🔍 Dolphin配置验证工具")
        print("=" * 50)
        
        checks = [
            ("Python版本", lambda: self.check_python_version()),
            ("依赖包", lambda: self.check_dependencies()),
            ("GPU支持", lambda: self.check_gpu_support()),
            ("目录结构", lambda: self.check_directories()),
            ("配置文件", lambda: self.check_config_file(config_path)),
            ("模型文件", lambda: self.check_model_files(config_path)),
            ("示例文件", lambda: self.check_demo_files()),
            ("模型加载", lambda: self.test_model_loading(config_path)),
        ]
        
        passed = 0
        total = len(checks)
        
        for check_name, check_func in checks:
            try:
                if check_func():
                    passed += 1
            except Exception as e:
                self.log_issue(f"{check_name}检查时发生异常: {e}")
        
        print("\n" + "=" * 50)
        print("📊 验证结果汇总")
        print(f"总检查项: {total}")
        print(f"通过项: {passed}")
        print(f"问题数: {len(self.issues)}")
        print(f"警告数: {len(self.warnings)}")
        
        if self.issues:
            print("\n❌ 发现的问题:")
            for i, issue in enumerate(self.issues, 1):
                print(f"  {i}. {issue}")
        
        if self.warnings:
            print("\n⚠️  警告信息:")
            for i, warning in enumerate(self.warnings, 1):
                print(f"  {i}. {warning}")
        
        if len(self.issues) == 0:
            print("\n🎉 配置验证通过！系统可以正常运行")
        else:
            print(f"\n⚠️  发现 {len(self.issues)} 个问题需要解决")
        
        return {
            "total_checks": total,
            "passed_checks": passed,
            "issues": self.issues,
            "warnings": self.warnings,
            "success": len(self.issues) == 0
        }


def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="Dolphin配置验证工具")
    parser.add_argument("--config", default="config/Dolphin.yaml", help="配置文件路径")
    parser.add_argument("--json", help="将结果保存为JSON文件")
    
    args = parser.parse_args()
    
    validator = ConfigValidator()
    result = validator.run_validation(args.config)
    
    if args.json:
        import json
        with open(args.json, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        print(f"\n📄 验证结果已保存到: {args.json}")
    
    # 返回适当的退出码
    exit_code = 0 if result["success"] else 1
    sys.exit(exit_code)


if __name__ == "__main__":
    main()